package utils

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
	"path/filepath"
)

type RSAUtil struct {
	PrivateKey *rsa.PrivateKey
	PublicKey  *rsa.PublicKey
}

func NewRSAUtil() *RSAUtil {
	return &RSAUtil{}
}

func (r *RSAUtil) GenerateKeyPair() error {
	// 创建cert目录（如果不存在）
	if err := os.MkdirAll("cert", 0755); err != nil {
		return err
	}

	// 生成密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	// 保存私钥
	privKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privKeyBytes,
	})
	err = os.WriteFile(filepath.Join("cert", "private.pem"), privKeyPEM, 0600)
	if err != nil {
		return err
	}

	// 保存公钥
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}
	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})
	return os.WriteFile(filepath.Join("cert", "public.pem"), pubKeyPEM, 0644)
}

func (r *RSAUtil) LoadKeyPair() error {
	// 检查cert目录是否存在
	if _, err := os.Stat("cert"); os.IsNotExist(err) {
		// 如果目录不存在，创建它
		if err := os.MkdirAll("cert", 0755); err != nil {
			return fmt.Errorf("创建cert目录失败: %v", err)
		}
	}

	// 检查密钥文件是否存在
	if _, err := os.Stat(filepath.Join("cert", "private.pem")); os.IsNotExist(err) {
		// 如果密钥不存在，生成新的密钥对
		if err := r.GenerateKeyPair(); err != nil {
			return fmt.Errorf("生成密钥对失败: %v", err)
		}
	}

	// 读取私钥
	privKeyData, err := os.ReadFile(filepath.Join("cert", "private.pem"))
	if err != nil {
		return fmt.Errorf("读取私钥文件失败: %v", err)
	}
	privBlock, _ := pem.Decode(privKeyData)
	if privBlock == nil {
		return errors.New("解析私钥失败")
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(privBlock.Bytes)
	if err != nil {
		return fmt.Errorf("解析私钥数据失败: %v", err)
	}
	r.PrivateKey = privateKey
	r.PublicKey = &privateKey.PublicKey
	return nil
}

func (r *RSAUtil) Encrypt(data []byte) ([]byte, error) {
	if r.PublicKey == nil {
		return nil, errors.New("public key is not loaded")
	}
	return rsa.EncryptPKCS1v15(rand.Reader, r.PublicKey, data)
}

func (r *RSAUtil) Decrypt(ciphertext []byte) ([]byte, error) {
	if r.PrivateKey == nil {
		return nil, errors.New("private key is not loaded")
	}
	return rsa.DecryptPKCS1v15(rand.Reader, r.PrivateKey, ciphertext)
}
